fix: avoid dictionary key overflow in group value output#23405
fix: avoid dictionary key overflow in group value output#23405StealthEyeLLC wants to merge 29 commits into
Conversation
Rich-T-kid
left a comment
Rich-T-kid
left a comment
There was a problem hiding this comment.
thank you for the PR @StealthEyeLLC, I left some of my thoughts.
| DataType::Dictionary(_, value_type) => DataType::Dictionary( | ||
| Box::new(DataType::UInt64), | ||
| Box::new(group_value_output_data_type(value_type)), | ||
| ), |
There was a problem hiding this comment.
🤔 I don't think making every dictionary key type u64 is the best option here. changing a key from i8 to u64 just because a group by had 128 groups is overkill in my opinion.
| let output_schema = group_value_output_schema(schema.as_ref()); | ||
| Ok(Self { | ||
| schema, | ||
| schema: output_schema, |
There was a problem hiding this comment.
The root of the issue is here, we don't know how many groups we will end up with until we've called emit(). to resolve #23127 I think we need to allow for schemas to be somewhat dynamic. In the sense that a dictionary's key type can group but no shrink.
|
Thanks for the review, @Rich-T-kid. I reworked the implementation around dynamic, monotonic key promotion rather than widening every dictionary key to UInt64. The current behavior is:
I added coverage for the promotion boundary, signed keys, no-shrink behavior, partial-to-final aggregation, the legacy path, and a repartitioned SQL regression. The current source passes the targeted Rust tests, SQLLogicTest, formatting, and clippy with warnings denied. The Apache workflow runs are still waiting for maintainer approval. Could you please take another look and approve the workflows when convenient? |
2 participants
Which issue does this PR close?
Closes #23127.
Rationale for this change
GROUP BY over dictionary-encoded columns can produce more distinct group values than the key type of any individual input dictionary can represent.
For example, every input Dictionary(UInt8, Utf8) batch can be valid while the aggregate result contains more than 256 distinct groups. Re-encoding the complete result with UInt8 keys then fails with DictionaryKeyOverflowError.
What changes are included in this PR?
This changes dictionary group-value output to promote key types only when required:
Partial aggregate output still crosses fixed-schema repartition and coalescing boundaries. To keep that transport stable, partial output is materialized once, sliced at the planned dictionary key capacity, and each slice is compactly re-encoded to the planned schema. Spill IPC similarly uses a stable internal dictionary width.
What regression coverage is included?
Are these changes tested?
Yes. The current source was validated with:
cargo fmt --check
cargo test -p datafusion-physical-plan dictionary_groups_keep_partial_schema_and_promote_final_output -- --nocapture
cargo test -p datafusion-physical-plan dict_uint8_utf8_promotes_and_does_not_shrink -- --nocapture
cargo test -p datafusion-physical-plan dictionary_group_key_promotes_runtime_schema -- --nocapture
cargo test --test sqllogictests -- dictionary_group_by_key_promotion
cargo clippy -p datafusion-physical-plan --lib --tests -- -D warnings
The upstream Apache workflows still require maintainer approval before they can run on this PR.